{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "acd755d3-f72d-4b91-8a88-bf3c3f165c21",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing 32 images of shape (512, 512)\n",
      "\n",
      "Sequential Case\n",
      "\tElapsed time: 378.52 ms\n",
      "\n",
      "Multithreaded Case\n",
      "\tNumber of concurrent workers: 48\n",
      "\tElapsed time: 65.67 ms\n",
      "\n",
      "Relative speedup with concurrent = 5.7642200975490745\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "Demo: Parallel processing accross images\n",
    "\n",
    "Multithreading can be used to run transforms on a set of images in parallel.\n",
    "This will give a net performance benefit if the images to be transformed are\n",
    "sufficiently large.\n",
    "\n",
    "This demo runs a multilevel wavelet decomposition on a list of 32 images,\n",
    "each of size (512, 512).  Computations are repeated sequentially and in\n",
    "parallel and the runtimes compared.\n",
    "\n",
    "In general, multithreading will be more beneficial for larger images and for\n",
    "wavelets with a larger filter size.\n",
    "\n",
    "One can also change ``ndim`` to 3 in the code below to use a set of 3D volumes\n",
    "instead.\n",
    "\"\"\"\n",
    "\n",
    "import time\n",
    "from functools import partial\n",
    "from multiprocessing import cpu_count\n",
    "\n",
    "try:\n",
    "    from concurrent import futures\n",
    "except ImportError:\n",
    "    raise ImportError(\n",
    "        \"This demo requires concurrent.futures.  It can be installed for \"\n",
    "        \"for python 2.x via:  pip install futures\")\n",
    "\n",
    "import numpy as np\n",
    "from numpy.testing import assert_array_equal\n",
    "\n",
    "import pywt\n",
    "\n",
    "# the test image\n",
    "cam = pywt.data.camera().astype(float)\n",
    "\n",
    "ndim = 2                   # dimension of images to transform (2 or 3)\n",
    "num_images = 32            # number of images to transform\n",
    "max_workers = cpu_count()  # max number of available threads\n",
    "nrepeat = 5                # averages used in the benchmark\n",
    "\n",
    "# create a list of num_images images\n",
    "if ndim == 2:\n",
    "    imgs = [cam, ] * num_images\n",
    "    wavelet = 'db8'\n",
    "elif ndim == 3:\n",
    "    # stack image along 3rd dimension to create a [512 x 512 x 16] 3D volume\n",
    "    im3 = np.concatenate([cam[:, :, np.newaxis], ]*16, axis=-1)\n",
    "    # create multiple copies of the volume\n",
    "    imgs = [im3, ] * num_images\n",
    "    wavelet = 'db1'\n",
    "else:\n",
    "    ValueError(\"Only 2D and 3D test cases implemented\")\n",
    "\n",
    "# define a function to apply to each image\n",
    "wavedecn_func = partial(pywt.wavedecn, wavelet=wavelet, mode='periodization',\n",
    "                        level=3)\n",
    "\n",
    "\n",
    "def concurrent_transforms(func, imgs, max_workers=None):\n",
    "    \"\"\"Call func on each img in imgs using a ThreadPoolExecutor.\"\"\"\n",
    "    executor = futures.ThreadPoolExecutor\n",
    "    if max_workers is None:\n",
    "        # default to as many workers as available cpus\n",
    "        max_workers = cpu_count()\n",
    "    results = []\n",
    "    with executor(max_workers=max_workers) as execute:\n",
    "        for result in execute.map(func, imgs):\n",
    "            results.append(result)\n",
    "    return results\n",
    "\n",
    "\n",
    "print(\"Processing {} images of shape {}\".format(len(imgs), imgs[0].shape))\n",
    "\n",
    "# Sequential computation via a list comprehension\n",
    "tstart = time.time()\n",
    "for n in range(nrepeat):\n",
    "    results = [wavedecn_func(img) for img in imgs]\n",
    "t = (time.time()-tstart)/nrepeat\n",
    "print(\"\\nSequential Case\")\n",
    "print(\"\\tElapsed time: {:0.2f} ms\".format(1000*t))\n",
    "\n",
    "\n",
    "# Concurrent computation via concurrent.futures\n",
    "tstart = time.time()\n",
    "for n in range(nrepeat):\n",
    "    results_concurrent = concurrent_transforms(wavedecn_func, imgs,\n",
    "                                               max_workers=max_workers)\n",
    "t2 = (time.time()-tstart)/nrepeat\n",
    "print(\"\\nMultithreaded Case\")\n",
    "print(\"\\tNumber of concurrent workers: {}\".format(max_workers))\n",
    "print(\"\\tElapsed time: {:0.2f} ms\".format(1000*t2))\n",
    "print(\"\\nRelative speedup with concurrent = {}\".format(t/t2))\n",
    "\n",
    "# check a couple of the coefficient arrays to verify matching results for\n",
    "# sequential and multithreaded computation\n",
    "assert_array_equal(results[-1][0],\n",
    "                   results_concurrent[-1][0])\n",
    "assert_array_equal(results[-1][1]['d' + 'a'*(ndim-1)],\n",
    "                   results_concurrent[-1][1]['d' + 'a'*(ndim-1)])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6665773f-9318-4643-82e0-cc59cd716338",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
